

from util.box_ops import box_cxcywh_to_xyxy
from .postprocesshoi import PostProcessHOI as posthoi
import torch
from torch import nn
import torch.nn.functional as F

from util.box_ops import box_cxcywh_to_xyxy
import numpy as np

class PostProcessHOI(posthoi):
    def __init__(self, args, mode='train', correct_mat = None, trans_type='detr'):
        super(PostProcessHOI, self).__init__(args, mode, correct_mat, trans_type)

    @torch.no_grad()
    def forward(self, outputs, target_sizes):
        sub_boxes, obj_boxes, obj_scores, obj_labels, hoi_scores  = \
            self.get_outputs(outputs, target_sizes)


        results = []
        for index in range(len(hoi_scores)):
            hs, os, ol, sb, ob = hoi_scores[index], obj_scores[index], obj_labels[index], sub_boxes[index], obj_boxes[index]
            sl = torch.full_like(ol, self.subject_category_id)
            l = torch.cat((sl, ol))
            b = torch.cat((sb, ob))
            
            ss = torch.full_like(os, 1.)
            s = torch.cat((ss, os))
            hoi_scores = hs # 100, 263
            
            ids = torch.arange(b.shape[0])

   
            if self.mode=='train':
                #   ['boxes', 'labels', 'score', 'sub_ids', 
                #   'obj_ids', 'verb_scores']
                results.append({'boxes':  np.clip(b.to('cpu'), 0, 1333),
                                'labels': l.to('cpu'),  
                                'score':  s.to('cpu')})
                results[-1].update({'sub_ids':     ids[:ids.shape[0] // 2],
                                    'obj_ids':     ids[ids.shape[0] // 2:],
                                    'hoi_scores':  hoi_scores.to('cpu'), })
#############################################  PNMS ###########################################################
            elif self.mode=='test':
                bboxes = [{'bbox':          np.clip(bbox, 0, 1333), 
                           'category_id':   label, 
                           'score':         score} 
                            for bbox, label, score in zip(b.to('cpu').numpy(), 
                                                          l.to('cpu').numpy(), 
                                                          s.to('cpu').numpy())]
                object_labels = ol.view(-1, 1).expand(-1, hoi_scores.shape[1])
                verb_labels = torch.arange(hoi_scores.shape[1], device=self.correct_mat.device).view(1, -1).expand(hoi_scores.shape[0], -1).to('cpu').numpy()
                subject_ids = np.tile(ids[:ids.shape[0] // 2].to('cpu').numpy(), (hoi_scores.shape[1], 1)).T
                object_ids = np.tile(ids[ids.shape[0] // 2:].to('cpu').numpy(), (hoi_scores.shape[1], 1)).T
                #   2900, 
                hoi_scores = hoi_scores.to('cpu').numpy().ravel()
                verb_labels = verb_labels.ravel()
                subject_ids = subject_ids.ravel()
                object_ids = object_ids.ravel()
                
                #   语义级别过滤矩阵
                masks = self.correct_mat[verb_labels.reshape(-1), 
                                         object_labels.reshape(-1)].view(hoi_scores.shape).to('cpu').numpy()
                hoi_scores *= masks
                hoi_scores[hoi_scores<self.query_thres] = 0.
                new_hoi_scores = hoi_scores[hoi_scores>0.] 
                verb_labels = verb_labels[hoi_scores>0.]
                object_ids =  object_ids[hoi_scores>0.]
                subject_ids = subject_ids[hoi_scores>0.]
                
                
                hois = [{'subject_id':  subject_id, 
                         'object_id':   object_id, 
                         'category_id': category_id, 
                         'score': score} for
                        subject_id, object_id, category_id, score in zip(subject_ids, object_ids, verb_labels, new_hoi_scores)]
                # hois.sort(key=lambda k: (k.get('score', 0)), reverse=True)
                # hois = hois[:self.max_hois]
                current_result = {'predictions': bboxes, 
                                  'hoi_prediction': hois}
                if self.use_nms_filter:
                    current_result = self.triplet_nms_filter(current_result)
                current_result['hoi_prediction'].sort(key=lambda k: (k.get('score', 0)), reverse=True)
                current_result['hoi_prediction'] = current_result['hoi_prediction'][:self.max_hois]
                results.append(current_result)
#############################################  PNMS ###########################################################

        return results


    @torch.no_grad()
    def get_outputs(self, outputs, target_sizes):
        assert target_sizes.shape[1] == 2
        obj_scores, obj_labels, topk_boxes = self.get_obj_cls(outputs['pred_obj_logits'], target_sizes)
        sub_boxes = self.get_sub_bbox(outputs['pred_sub_boxes'], target_sizes, topk_boxes)
        obj_boxes = self.get_obj_bbox(outputs['pred_obj_boxes'], target_sizes, topk_boxes)
        hoi_scores = self.get_hoi_score(outputs['pred_hoi_logits'], target_sizes, topk_boxes)

        return sub_boxes, obj_boxes, obj_scores, obj_labels, hoi_scores

    def get_hoi_score(self, pred_hoi_logits, target_sizes, topk_boxes, topK=20):
        out_hoi_logits = pred_hoi_logits
        #out_hoi_logits = outputs['pred_hoi_logits_cascade'].permute(1,0,2,3).flatten(1,2)
        hoi_scores = out_hoi_logits.sigmoid()
        N, num_query, num_class = hoi_scores.shape
        #   100, 29
        if self.trans_type=='detr':
            _, hoi_index  = torch.topk(-hoi_scores, num_class - topK, dim=-1)
            hoi_scores = hoi_scores.scatter(-1, hoi_index, 0)
        elif self.trans_type=='ddetr':
            hoi_scores = torch.gather(hoi_scores, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, num_class))
        return hoi_scores